"""
    测试图片像素输出
"""
import os

import numpy as np
import torch
from albumentations import Compose, Resize, Normalize
from albumentations.pytorch import ToTensorV2
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from itertools import islice

from evaluate.evaluate_pth import CrackTestDataset
from model.deepcrack_with_uu_Sp_3 import DeepCrack_uu_Sp_3
from trainer.trainer import DeepCrackTrainer
from config.config import Config as cfg  # 配置文件（包含模型参数、训练参数等）
# 设置CUDA可见设备（指定使用的GPU）
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu_id

mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
test_transform = Compose([
        # Resize(256, 256),
        # Normalize(mean=mean, std=std),
        ToTensorV2(),
    ])

# 创建测试集
test_dataset = CrackTestDataset(

    # win下用相对路径
    # image_path=r'data/validation/images',
    # mask_path=r'data/validation/label',

    # linux下使用绝对路径
    img_dir='/home/likang/DeepCrack_LK/evaluate/data/validation/images',
    mask_dir='/home/likang/DeepCrack_LK/evaluate/data/validation/label',
    transform=test_transform
)
test_loader = DataLoader(test_dataset, batch_size=4, shuffle=False)
# print(test_loader.dataset)
# torch.no_grad() 表示不跟踪梯度，可以提高运行速度

# ---------------------引入和调整模型------------------
# 预训练模型权重路径（注意：路径需根据实际位置修改）
pretrained_model = r'../../may_predict/predicted_model/DuSp3_CRKWH1001_Wight6_2_200_DuSp3_CRKWH1001_Wight6_2_200_epoch(199)_0000200_2025-05-08-13-25-51.pth'

# ==================== 2. 模型加载与配置 ====================
# 设置计算设备（优先使用GPU）
device = torch.device("cuda")

# 初始化模型架构
model = DeepCrack_uu_Sp_3(cfg)  # 使用配置文件参数构建模型

# 多GPU并行计算配置（device_ids指定使用的GPU索引）
model = torch.nn.DataParallel(model, device_ids=[0]).cuda()

# 初始化训练器并加载预训练权重
trainer = DeepCrackTrainer(model).cuda()
model.load_state_dict(
    trainer.saver.load(pretrained_model, multi_gpu=True),
    strict=True  # 严格匹配权重名称
)
model.eval()  # 设置为评估模式（关闭dropout等训练专用层）
with torch.no_grad():
    # 切片前num张
    num = 5
    for i, (images, masks) in islice(enumerate(test_loader),num):
        # 预处理的原图shape
        # print(images.shape)  # torch.Size([4, 3, 512, 512])  4是批次大小 DataLoader里面手动设定的
        # print(images)


        # 反标准化
        images_ = images * torch.tensor(std).view(3, 1, 1) + torch.tensor(mean).view(3, 1, 1)

        # 获取当前批次对应的文件名
        filenames = test_dataset.get_image_filenames(i * test_loader.batch_size, (i + 1) * test_loader.batch_size)

        # 展示images的图像
        print("原图shape:")
        print(images_[0].shape) #torch.Size([3, 256, 256])
        # transpose(1, 2, 0) 将(C,H,W) 转换为(H,W,C) 因为matplotlib的imshow只能处理(H,W,C)
        plt.imshow(images_[0].cpu().numpy().transpose(1, 2, 0))
        plt.title(filenames[0])# 标题
        plt.show()


        print("标签shape:")
        print(masks.shape)#torch.Size([4, 512, 512])
        print(masks[0].shape) #torch.Size([256, 256])

        # print(masks)
        # 展示masks的图像    这里为什么取masks[0] 因为masks的定义是4通道，总共有4张图，所以取masks[0]第一张
        plt.imshow(masks[0].cpu().numpy())
        plt.title(filenames[0])  # 标题
        plt.show()

        images = images_[0].cpu().numpy()
        images = images.astype(np.float32) / 255.0
        test_data = torch.from_numpy(images).float().unsqueeze(0).to(device)
        pred = model(test_data)
        # print(torch.sigmoid(pred[0].cpu().squeeze()))
        pred = torch.sigmoid(pred[0].cpu().squeeze())
        # print(pred[0].max())
        # print("多尺度预测 的第一个全尺度预测结果的shape:")
        # print(pred[0].shape) #torch.Size([4, 1, 512, 512])
        # print("多尺度预测 的第一个全尺度预测结果的第一个通道 即4张图的第一张图:")
        # print(pred[0][0].shape) #torch.Size([1, 512, 512])


        showMode = "color"  # mask
        if showMode == "color":
            # 原图的预测结果   彩色展示
            # pred_sigmoid = torch.sigmoid(pred[0][0].cpu())
            print(pred[0].shape)
            print(pred[0][0].shape)
            plt.imshow(pred.cpu().numpy())
            plt.show()
        elif showMode == "mask":
            save_pred = pred[0][0].squeeze()
            pree = np.power(save_pred.cpu().data.numpy(), 0.7)  # 指数0.7为经验值
            plt.imshow(pree, cmap='gray')
            plt.title(filenames[0])  # 标题
            plt.show()
